from tortoise import Model, fields
from tortoise.contrib.pydantic import pydantic_model_creator

class Stock(Model):
    id = fields.IntField(pk=True)
    symbol = fields.CharField(max_length=20, unique=True)
    name = fields.CharField(max_length=100)
    created_at = fields.DatetimeField(auto_now_add=True)
    
    class Meta:
        table = "stocks"

class StockPrice(Model):
    id = fields.IntField(pk=True)
    symbol = fields.CharField(max_length=20)
    date = fields.DateField()
    open_price = fields.DecimalField(max_digits=10, decimal_places=2)
    high_price = fields.DecimalField(max_digits=10, decimal_places=2)
    low_price = fields.DecimalField(max_digits=10, decimal_places=2)
    close_price = fields.DecimalField(max_digits=10, decimal_places=2)
    volume = fields.DecimalField(max_digits=10, decimal_places=2)
    amount = fields.DecimalField(max_digits=15, decimal_places=2)
    amplitude = fields.DecimalField(max_digits=5, decimal_places=2)
    pct_change = fields.DecimalField(max_digits=5, decimal_places=2)
    change_amount = fields.DecimalField(max_digits=10, decimal_places=2)
    turnover = fields.DecimalField(max_digits=5, decimal_places=2)
    
    class Meta:
        table = "stock_prices"
        indexes = [("symbol", "date")]

class NewsArticle(Model):
    id = fields.IntField(pk=True)
    stock = fields.ForeignKeyField("models.Stock", related_name="news")
    title = fields.CharField(max_length=200)
    content = fields.TextField()
    url = fields.CharField(max_length=500)
    published_at = fields.DatetimeField()
    sentiment_score = fields.FloatField(null=True)
    
    class Meta:
        table = "news_articles"

class TradingSignal(Model):
    id = fields.IntField(pk=True)
    stock = fields.ForeignKeyField("models.Stock", related_name="signals")
    signal_type = fields.CharField(max_length=50)  # 'buy', 'sell', 'hold'
    confidence = fields.FloatField()
    reasoning = fields.TextField()
    created_at = fields.DatetimeField(auto_now_add=True)
    
    class Meta:
        table = "trading_signals"

class Portfolio(Model):
    id = fields.IntField(pk=True)
    cash = fields.DecimalField(max_digits=15, decimal_places=2)
    created_at = fields.DatetimeField(auto_now_add=True)
    updated_at = fields.DatetimeField(auto_now=True)
    
    class Meta:
        table = "portfolios"

class PortfolioPosition(Model):
    id = fields.IntField(pk=True)
    portfolio = fields.ForeignKeyField("models.Portfolio", related_name="positions")
    stock = fields.ForeignKeyField("models.Stock", related_name="portfolio_positions")
    quantity = fields.IntField()
    average_price = fields.DecimalField(max_digits=10, decimal_places=2)
    created_at = fields.DatetimeField(auto_now_add=True)
    updated_at = fields.DatetimeField(auto_now=True)
    
    class Meta:
        table = "portfolio_positions"

# 创建Pydantic模型用于API响应
Stock_Pydantic = pydantic_model_creator(Stock, name="Stock")
StockPrice_Pydantic = pydantic_model_creator(StockPrice, name="StockPrice")
NewsArticle_Pydantic = pydantic_model_creator(NewsArticle, name="NewsArticle")
TradingSignal_Pydantic = pydantic_model_creator(TradingSignal, name="TradingSignal")
Portfolio_Pydantic = pydantic_model_creator(Portfolio, name="Portfolio")
PortfolioPosition_Pydantic = pydantic_model_creator(PortfolioPosition, name="PortfolioPosition")
